package br.unifor.g2cl;

import java.io.IOException;
import java.net.SocketAddress;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.KeySpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;

import net.sf.jgcs.Annotation;
import net.sf.jgcs.ControlSession;
import net.sf.jgcs.Message;
import net.sf.jgcs.Service;
import net.sf.jgcs.UnsupportedServiceException;
import br.unifor.g2cl.annotation.PostConstruct;

public class CryptoDataSession extends DataSessionAdapter {
	
	private byte[] ENCRYPTED = new byte[]{-56};
	private byte[] NOT_ENCRYPTED = new byte[]{-57};
	
	String passPhrase = "GenericGroupCommunicationLayer";

	public CryptoDataSession(IMarshalDataSession data, ControlSession control) {
		super(data);
	}
	
	DesEncrypter desEncrypt;
	
	private DataSessionWorker worker;
	
	public void setWorker(DataSessionWorker worker){
		this.worker = worker;
	}
	
	@PostConstruct
	public void init(){
		if(worker == null){
			desEncrypt = new DesEncrypter(passPhrase);	
		}
	}
	
	@Override
	public void multicast(Message msg, Service service, Object cookie,
			Annotation... annotation) throws IOException,
			UnsupportedServiceException {
		
		G2CLMessage msgEnc = (G2CLMessage)createMessage();
		msgEnc.setSenderAddress(msg.getSenderAddress());
		byte[] payload = msg.getPayload();
		
		byte[] encPayload = encrypt(payload);
		if(encPayload == null){
			msgEnc.setPayload(payload);
			msgEnc.addHeader(NOT_ENCRYPTED);
		}
		else{
			msgEnc.setPayload(encPayload);
			msgEnc.addHeader(ENCRYPTED);
		}
		
				
		super.multicast(msgEnc, service, cookie, annotation);
		
	}
	
	@Override
	public void send(Message msg, Service service, Object cookie,
			SocketAddress destination, Annotation... annotation)
			throws IOException, UnsupportedServiceException {		

		G2CLMessage msgEnc = (G2CLMessage)createMessage();
		msgEnc.setSenderAddress(msg.getSenderAddress());
		byte[] payload = msg.getPayload();
		
		byte[] encPayload = encrypt(payload);
		if(encPayload == null){
			msgEnc.setPayload(payload);
			msgEnc.addHeader(NOT_ENCRYPTED);
		}
		else{
			msgEnc.setPayload(encPayload);
			msgEnc.addHeader(ENCRYPTED);
		}
				
		super.send(msg, service, cookie, destination, annotation);
	}
	
	@Override
	public Object onMessage(Message msg) {

		G2CLMessage message = (G2CLMessage)msg;
		byte[] header = message.removeHeader();
		
		if(header[0] == ENCRYPTED[0]){
			byte[] payload = message.getPayload();		
			message.setPayload(decrypt(payload));			
		}		
		
		return super.onMessage(message);
	}
	
	public byte[] encrypt(byte[] data){
		
		if(worker != null){
			return worker.doIt(data);
		}
		
		return desEncrypt.encrypt(data);
		
	}
	
	public byte[] decrypt(byte[] data){
		if(worker != null){
			return worker.undoIt(data);
		}
		
		return desEncrypt.decrypt(data);
		
	}
	
    class DesEncrypter {
        Cipher ecipher;
        Cipher dcipher;
    
        // 8-byte Salt
        byte[] salt = {
            (byte)0xA9, (byte)0x9B, (byte)0xC8, (byte)0x32,
            (byte)0x56, (byte)0x35, (byte)0xE3, (byte)0x03
        };
    
        // Iteration count
        int iterationCount = 19;
    
        DesEncrypter(String passPhrase) {
            try {
                // Create the key
                KeySpec keySpec = new PBEKeySpec(passPhrase.toCharArray(), salt, iterationCount);
                SecretKey key = SecretKeyFactory.getInstance(
                    "PBEWithMD5AndDES").generateSecret(keySpec);
                ecipher = Cipher.getInstance(key.getAlgorithm());
                dcipher = Cipher.getInstance(key.getAlgorithm());
    
                // Prepare the parameter to the ciphers
                AlgorithmParameterSpec paramSpec = new PBEParameterSpec(salt, iterationCount);
    
                // Create the ciphers
                ecipher.init(Cipher.ENCRYPT_MODE, key, paramSpec);
                dcipher.init(Cipher.DECRYPT_MODE, key, paramSpec);
            } catch (Exception e) {
            	e.printStackTrace();
            }
        }
    
        public byte[] encrypt(byte[] dataToEncrypt) {
            try {
                // Encrypt
                byte[] enc = ecipher.doFinal(dataToEncrypt);
    
                return enc;
            } catch (Exception e) {
            	e.printStackTrace();
            }
            return null;
        }
    
        public byte[] decrypt(byte[] encryptedData) {
            try {
                // Decrypt
                return dcipher.doFinal(encryptedData);
            } catch (Exception e) {
            	e.printStackTrace();
            }
            return null;
        }        
    }

	public String getPassPhrase() {
		return passPhrase;
	}

	public void setPassPhrase(String passPhrase) {
		this.passPhrase = passPhrase;
	}
}
